package com.rent.message.interceptor;

import com.rent.common.constant.CommonConstant;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageChannel;
import org.springframework.messaging.simp.stomp.StompCommand;
import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
import org.springframework.messaging.support.ChannelInterceptor;
import org.springframework.messaging.support.MessageHeaderAccessor;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

import java.security.Principal;
import java.util.Map;

@Component
@Slf4j
public class StompChannelAuthInterceptor implements ChannelInterceptor {

    @Autowired
    private RedisTemplate<String, Object> redisTemplate;

    @Override
    public Message<?> preSend(Message<?> message, MessageChannel channel) {
        StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
        
        if (accessor == null) {
            return message;
        }

        String sessionId = accessor.getSessionId();
        StompCommand command = accessor.getCommand();

        // 对连接和订阅命令进行验证
        if (StompCommand.CONNECT.equals(command) || StompCommand.SUBSCRIBE.equals(command)) {
            // 从native headers获取token
            Map<String, Object> sessionAttributes = accessor.getSessionAttributes();
            String userId = sessionAttributes != null ? (String) sessionAttributes.get("userId") : null;
            
            if (!StringUtils.hasText(userId)) {
                log.warn("STOMP命令拒绝: 未认证的用户, sessionId={}, command={}", sessionId, command);
                return null; // 拒绝消息
            }

            // 设置用户principal，用于@SendToUser
            accessor.setUser(new Principal() {
                @Override
                public String getName() {
                    return userId;
                }
            });
        }

        // 特别保护注册和发送消息的端点
        if (StompCommand.SEND.equals(command)) {
            String destination = accessor.getDestination();
            if (destination != null && 
                (destination.startsWith("/app/message/register") || 
                 destination.startsWith("/app/chat/send"))) {
                
                Map<String, Object> sessionAttributes = accessor.getSessionAttributes();
                String userId = sessionAttributes != null ? (String) sessionAttributes.get("userId") : null;
                
                if (!StringUtils.hasText(userId)) {
                    log.warn("消息发送拒绝: 未认证的用户尝试发送消息, sessionId={}, destination={}", 
                            sessionId, destination);
                    return null;
                }
            }
        }

        return message;
    }
}